{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "850de06d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Node1] 当前 step_count: 0\n",
      "[Node2] 当前 step_count: 1\n",
      "\n",
      " 最终输出:\n",
      "{'messages': ['Initial input', 'Hello from node1 at step 1', 'Goodbye from node2 at step 2'], 'step_count': 2}\n",
      "\n",
      " 状态变更历史（快照）:\n",
      "StateSnapshot(values={'messages': ['Initial input', 'Hello from node1 at step 1', 'Goodbye from node2 at step 2'], 'step_count': 2}, next=(), config={'configurable': {'thread_id': '123', 'checkpoint_ns': '', 'checkpoint_id': '1f090bb7-f882-6a9c-8002-c61205ac2bf9'}}, metadata={'source': 'loop', 'step': 2, 'parents': {}}, created_at='2025-09-13T16:05:42.139970+00:00', parent_config={'configurable': {'thread_id': '123', 'checkpoint_ns': '', 'checkpoint_id': '1f090bb7-f882-6a9b-8001-220d0d477fdb'}}, tasks=(), interrupts=())\n",
      "StateSnapshot(values={'messages': ['Initial input', 'Hello from node1 at step 1'], 'step_count': 1}, next=('node2',), config={'configurable': {'thread_id': '123', 'checkpoint_ns': '', 'checkpoint_id': '1f090bb7-f882-6a9b-8001-220d0d477fdb'}}, metadata={'source': 'loop', 'step': 1, 'parents': {}}, created_at='2025-09-13T16:05:42.139970+00:00', parent_config={'configurable': {'thread_id': '123', 'checkpoint_ns': '', 'checkpoint_id': '1f090bb7-f880-62b9-8000-c8e66e66edff'}}, tasks=(PregelTask(id='3c707d99-0376-3a98-d27d-5d98b7c7baae', name='node2', path=('__pregel_pull', 'node2'), error=None, interrupts=(), state=None, result={'messages': ['Goodbye from node2 at step 2'], 'step_count': 2}),), interrupts=())\n",
      "StateSnapshot(values={'messages': ['Initial input'], 'step_count': 0}, next=('node1',), config={'configurable': {'thread_id': '123', 'checkpoint_ns': '', 'checkpoint_id': '1f090bb7-f880-62b9-8000-c8e66e66edff'}}, metadata={'source': 'loop', 'step': 0, 'parents': {}}, created_at='2025-09-13T16:05:42.138949+00:00', parent_config={'configurable': {'thread_id': '123', 'checkpoint_ns': '', 'checkpoint_id': '1f090bb7-f87d-6b68-bfff-88691867fd35'}}, tasks=(PregelTask(id='a62460d6-729d-42d1-473c-8a5e5724df02', name='node1', path=('__pregel_pull', 'node1'), error=None, interrupts=(), state=None, result={'messages': ['Hello from node1 at step 1'], 'step_count': 1}),), interrupts=())\n",
      "StateSnapshot(values={'messages': []}, next=('__start__',), config={'configurable': {'thread_id': '123', 'checkpoint_ns': '', 'checkpoint_id': '1f090bb7-f87d-6b68-bfff-88691867fd35'}}, metadata={'source': 'input', 'step': -1, 'parents': {}}, created_at='2025-09-13T16:05:42.137943+00:00', parent_config=None, tasks=(PregelTask(id='ffeb084c-5a8e-a45c-25f2-a921c32ccc71', name='__start__', path=('__pregel_pull', '__start__'), error=None, interrupts=(), state=None, result={'messages': ['Initial input'], 'step_count': 0}),), interrupts=())\n"
     ]
    }
   ],
   "source": [
    "from typing import Annotated, Sequence\n",
    "from pydantic import BaseModel, Field\n",
    "import operator\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.checkpoint.memory import MemorySaver  # ←← \n",
    "\n",
    "# === 1. 定义状态 Schema (Pydantic) ===\n",
    "class AgentState(BaseModel):\n",
    "    messages: Annotated[Sequence[str], operator.add] = Field(default_factory=list)\n",
    "    step_count: int = 0\n",
    "\n",
    "\n",
    "# === 2. 节点函数 ===\n",
    "def node1(state: AgentState) -> dict:\n",
    "    print(f\"[Node1] 当前 step_count: {state.step_count}\")\n",
    "    new_message = f\"Hello from node1 at step {state.step_count + 1}\"\n",
    "    return {\n",
    "        \"messages\": [new_message],\n",
    "        \"step_count\": state.step_count + 1,\n",
    "    }\n",
    "\n",
    "\n",
    "def node2(state: AgentState) -> dict:\n",
    "    print(f\"[Node2] 当前 step_count: {state.step_count}\")\n",
    "    new_message = f\"Goodbye from node2 at step {state.step_count + 1}\"\n",
    "    return {\n",
    "        \"messages\": [new_message],\n",
    "        \"step_count\": state.step_count + 1,\n",
    "    }\n",
    "\n",
    "\n",
    "# === 3. 构建图 ===\n",
    "builder = StateGraph(AgentState)\n",
    "\n",
    "builder.add_node(\"node1\", node1)\n",
    "builder.add_node(\"node2\", node2)\n",
    "\n",
    "builder.add_edge(START, \"node1\")\n",
    "builder.add_edge(\"node1\", \"node2\")\n",
    "builder.add_edge(\"node2\", END)\n",
    "\n",
    "# 编译时必须传入 checkpointer\n",
    "memory = MemorySaver()\n",
    "graph = builder.compile(checkpointer=memory)  # ←← 启用 checkpointer\n",
    "\n",
    "# === 4. 执行图并传入 config（包含 thread_id）===\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = graph.invoke(\n",
    "    {\"messages\": [\"Initial input\"], \"step_count\": 0},\n",
    "    config=config\n",
    ")\n",
    "\n",
    "print(\"\\n 最终输出:\")\n",
    "print(result)\n",
    "\n",
    "# === 5. 获取状态历史（快照）===\n",
    "print(\"\\n 状态变更历史（快照）:\")\n",
    "state_history = list(graph.get_state_history(config))\n",
    "for state in state_history:\n",
    "    print(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48ec82b6",
   "metadata": {},
   "source": [
    "可以根据检查点 ID 回到任意一个状态，并从那里重新启动 Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "494d700b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Node1] 当前 step_count: 0\n",
      "[Node2] 当前 step_count: 1\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'messages': ['Initial input',\n",
       "  'Hello from node1 at step 1',\n",
       "  'Initial input',\n",
       "  'Hello from node1 at step 1',\n",
       "  'Goodbye from node2 at step 2'],\n",
       " 'step_count': 2}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = {\"configurable\":{'thread_id': '123',  'checkpoint_id': '1f090bb7-f882-6a9b-8001-220d0d477fdb'}}\n",
    "\n",
    "checkpoint_snapshot = graph.get_state(config)\n",
    "\n",
    "graph.invoke({\"messages\": [\"Initial input\"], \"step_count\": 0}, config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "38e106d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qU redis aioredis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1e26c7cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "查询订单: 订单号不合法，请输入10到12位数字。\n",
      "Messages: [HumanMessage(content='abc123', additional_kwargs={}, response_metadata={}, id='3439bc1e-ac46-4499-838f-1a3c9aa2e227'), HumanMessage(content='订单号不合法，请输入10到12位数字。', additional_kwargs={}, response_metadata={}, id='a79cf0d3-3c31-4cd7-8e69-9d98500e8311'), HumanMessage(content='订单状态: 已发货', additional_kwargs={}, response_metadata={}, id='c4cf2943-f121-4bee-9639-fc9ed7e01919')]\n",
      "Retry count: 1\n",
      "---\n",
      "Messages: [HumanMessage(content='abc123', additional_kwargs={}, response_metadata={}, id='3439bc1e-ac46-4499-838f-1a3c9aa2e227'), HumanMessage(content='订单号不合法，请输入10到12位数字。', additional_kwargs={}, response_metadata={}, id='a79cf0d3-3c31-4cd7-8e69-9d98500e8311')]\n",
      "Retry count: 1\n",
      "---\n",
      "Messages: [HumanMessage(content='abc123', additional_kwargs={}, response_metadata={}, id='3439bc1e-ac46-4499-838f-1a3c9aa2e227'), HumanMessage(content='订单号不合法，请输入10到12位数字。', additional_kwargs={}, response_metadata={}, id='a79cf0d3-3c31-4cd7-8e69-9d98500e8311')]\n",
      "Retry count: 1\n",
      "---\n",
      "Messages: [HumanMessage(content='abc123', additional_kwargs={}, response_metadata={}, id='3439bc1e-ac46-4499-838f-1a3c9aa2e227')]\n",
      "Retry count: 0\n",
      "---\n",
      "Messages: []\n",
      "Retry count: 0\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "from typing import Annotated, Sequence, Literal\n",
    "from typing_extensions import TypedDict\n",
    "from langgraph.graph import StateGraph, START, END, add_messages\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "\n",
    "# 1. 定义状态\n",
    "class GraphState(TypedDict):\n",
    "    messages: Annotated[Sequence, add_messages]\n",
    "    retry_count: int\n",
    "\n",
    "\n",
    "# 2. 输入校验函数\n",
    "def is_valid_order_id(text: str) -> bool:\n",
    "    return bool(re.fullmatch(r\"\\d{10,12}\", text))\n",
    "\n",
    "\n",
    "def validate_input(state: GraphState) -> Literal[\"valid\", \"invalid\"]:\n",
    "    last_message = state[\"messages\"][-1]\n",
    "    user_input = last_message.content.strip()\n",
    "    if is_valid_order_id(user_input):\n",
    "        return \"valid\"\n",
    "    return \"invalid\"\n",
    "\n",
    "\n",
    "# 3. 节点函数\n",
    "def receive_input(state: GraphState):\n",
    "    user_input = state[\"messages\"][-1].content.strip()\n",
    "    return {\"order_id\": user_input}\n",
    "\n",
    "\n",
    "def query_order(state: GraphState):\n",
    "    order_id = state[\"messages\"][-1].content.strip()\n",
    "    print(f\"查询订单: {order_id}\")\n",
    "    return {\n",
    "        \"messages\": [\"订单状态: 已发货\"],\n",
    "        \"status\": \"success\"\n",
    "    }\n",
    "\n",
    "\n",
    "def handle_invalid(state: GraphState):\n",
    "    retry = state.get(\"retry_count\", 0)\n",
    "    if retry >= 2:\n",
    "        reply = \"输入错误次数过多，会话结束。\"\n",
    "        return {\"messages\": [reply], \"retry_count\": retry + 1}\n",
    "    \n",
    "    reply = \"订单号不合法，请输入10到12位数字。\"\n",
    "    return {\"messages\": [reply], \"retry_count\": retry + 1}\n",
    "\n",
    "\n",
    "# 4. 构建图\n",
    "builder = StateGraph(GraphState)\n",
    "\n",
    "builder.add_node(\"receive_input\", receive_input)\n",
    "builder.add_node(\"query_order\", query_order)\n",
    "builder.add_node(\"handle_invalid\", handle_invalid)\n",
    "\n",
    "# 条件跳转\n",
    "builder.add_conditional_edges(\n",
    "    START,\n",
    "    validate_input,\n",
    "    {\n",
    "        \"valid\": \"query_order\",\n",
    "        \"invalid\": \"handle_invalid\"\n",
    "    }\n",
    ")\n",
    "\n",
    "# 如果无效，最多重试2次\n",
    "builder.add_conditional_edges(\n",
    "    \"handle_invalid\",\n",
    "    lambda s: \"receive_input\" if s[\"retry_count\"] < 2 else END\n",
    ")\n",
    "\n",
    "builder.add_edge(\"receive_input\", \"query_order\")\n",
    "builder.add_edge(\"query_order\", END)\n",
    "\n",
    "# 编译图，使用 MemorySaver 记录状态历史\n",
    "memory_saver = MemorySaver()\n",
    "app = builder.compile(checkpointer=memory_saver)\n",
    "\n",
    "\n",
    "# 5. 测试运行\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "\n",
    "# 第一次：非法输入\n",
    "app.invoke(\n",
    "    {\"messages\": [\"abc123\"]},\n",
    "    config=config\n",
    ")\n",
    "\n",
    "# 查看状态历史\n",
    "history = app.get_state_history(config)\n",
    "for snapshot in history:\n",
    "    print(\"Messages:\", snapshot.values[\"messages\"])\n",
    "    print(\"Retry count:\", snapshot.values.get(\"retry_count\", 0))\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1e29866b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, Sequence\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import StateGraph, START, END, add_messages\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, AIMessage\n",
    "from langgraph.checkpoint.memory import MemorySaver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9ba18aba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "class GraphState(TypedDict):\n",
    "    messages: Annotated[Sequence, add_messages]\n",
    "    current_model: str\n",
    "\n",
    "# 主模型（故意设错 API Key 来模拟调用失败）\n",
    "primary_model = ChatOpenAI(\n",
    "    model=\"gpt-4\",\n",
    "    api_key=\"invalid_key_for_test\",  # 强制报错\n",
    "    temperature=0.7,\n",
    "    timeout=5.0\n",
    ")\n",
    "\n",
    "# 备用模型（使用你真实的 API Key）\n",
    "backup_model = ChatOpenAI(\n",
    "    model=\"gpt-3.5-turbo\",\n",
    "    api_key=os.getenv(\"OPENAI_API_KEY\"),  # 替换为你的有效 key\n",
    "    temperature=0.7,\n",
    "    timeout=10.0\n",
    ") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a7220ad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def call_llm(state):\n",
    "    messages = state[\"messages\"]\n",
    "    used_model = \"\"\n",
    "    response_message = None\n",
    "\n",
    "    # 尝试主模型\n",
    "    try:\n",
    "        print(\"尝试调用主模型 (gpt-4)...\")\n",
    "        response = primary_model.invoke(messages)\n",
    "        response_message = response\n",
    "        used_model = \"gpt-4\"\n",
    "    except Exception as e:\n",
    "        print(f\"主模型失败: {type(e).__name__}: {str(e)}\")\n",
    "        print(\"切换到备用模型 (gpt-3.5-turbo)\")\n",
    "\n",
    "        # 尝试备用模型\n",
    "        try:\n",
    "            response = backup_model.invoke(messages)\n",
    "            response_message = response\n",
    "            used_model = \"gpt-3.5-turbo\"\n",
    "        except Exception as e2:\n",
    "            print(f\"备用模型也失败: {str(e2)}\")\n",
    "            return {\n",
    "                \"messages\": [AIMessage(content=\"服务暂时不可用，请稍后再试。\")],\n",
    "                \"current_model\": \"none\"\n",
    "            }\n",
    "\n",
    "    return {\n",
    "        \"messages\": [AIMessage(content=response_message.content)],\n",
    "        \"current_model\": used_model\n",
    "    }\n",
    "\n",
    "builder = StateGraph(GraphState)\n",
    "builder.add_node(\"call_llm\", call_llm)\n",
    "builder.add_edge(START, \"call_llm\")\n",
    "builder.add_edge(\"call_llm\", END)\n",
    "\n",
    "# 启用检查点，用于后续查看状态历史\n",
    "checkpointer = MemorySaver()\n",
    "app = builder.compile(checkpointer=checkpointer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9c5e7eb6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "尝试调用主模型 (gpt-4)...\n",
      "主模型失败: AuthenticationError: Error code: 401 - {'error': {'message': '无效的令牌 (request id: 2025091400414229283963658964709)', 'type': 'v_api_error'}}\n",
      "切换到备用模型 (gpt-3.5-turbo)\n"
     ]
    }
   ],
   "source": [
    "# 测试\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"test_001\"}}\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [HumanMessage(content=\"解释一下什么是机器学习\")],\n",
    "    \"current_model\": \"\"\n",
    "}\n",
    "\n",
    "result = app.invoke(inputs, config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d60136cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "最终使用的模型: gpt-3.5-turbo\n",
      "AI 回复:\n",
      "机器学习是一种人工智能的分支，其目的是让计算机系统通过学习数据和模式来自动改进和适应，而不需要明确的编程指令。通过机器学习，计算机系统可以从大量的数据中识别模式和趋势，并利用这些信息来做出预测和决策。机器学习的应用范围非常广泛，包括语音识别、图像识别、自然语言处理、推荐系统等领域。在机器学习中，常用的算法包括监督学习、无监督学习和强化学习等。通过不断地训练和优化，机器学习算法可以不断提高性能和准确性，从而更好地应用于实际问题中。\n"
     ]
    }
   ],
   "source": [
    "print(\"最终使用的模型:\", result[\"current_model\"])\n",
    "\n",
    "for msg in result[\"messages\"]:\n",
    "    if isinstance(msg, AIMessage):\n",
    "        print(\"AI 回复:\")\n",
    "        print(msg.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MLOps",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
